import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist

from nanovllm.utils.context import get_context


class VocabParallelEmbedding(nn.Module):
    """
    这段代码实现了大语言模型（LLM）分布式训练 / 推理中的词汇表并行（Vocab Parallelism） 组件，包括 VocabParallelEmbedding（并行化词嵌入层）和 ParallelLMHead（并行化语言模型输出头）。核心目的是解决大词汇表（如百万级 token）导致的单卡内存不足问题 —— 通过将词汇表拆分为多份，由不同 GPU 分片存储和计算，再通过分布式通信汇总结果，是张量并行（Tensor Parallelism） 的关键子模块。
核心背景：为什么需要词汇表并行？
LLM 的词嵌入层（Embedding）和语言模型头（LMHead）的参数规模与词汇表大小（num_embeddings） 直接相关。例如：

若词汇表大小为 100 万，嵌入维度为 4096，则词嵌入层参数为 1e6 * 4096 = 4.096e9（约 16GB，float32）。
单卡难以容纳如此大的参数，因此将词汇表按 GPU 数量（tp_size）拆分为 tp_size 份，每个 GPU 仅存储 1/tp_size 的参数，大幅降低单卡内存占用。
一、VocabParallelEmbedding：并行化词嵌入层
负责将输入的 token_id 映射为嵌入向量，核心是分片存储嵌入权重 + 分布式汇总结果。

分布式参数依赖：dist.get_rank() 和 dist.get_world_size() 需基于 torch.distributed 初始化（如 dist.init_process_group(backend="nccl")），否则报错。
分片范围计算：例如 tp_size=4、num_embeddings=10000，则每个 GPU 存储 2500 个词汇的嵌入，rank0 负责 0-2499，rank1 负责 2500-4999，以此类推。
自定义权重加载：weight_loader 用于加载预训练权重时，仅切分并加载当前 GPU 负责的分片（避免全量权重加载导致内存爆炸）。

三、关键设计对比：VocabParallelEmbedding vs ParallelLMHead
模块	核心操作	分布式通信方式	输出结果用途
VocabParallelEmbedding	词嵌入 lookup	dist.all_reduce（所有 GPU 汇总）	模型输入嵌入向量
ParallelLMHead	线性变换（隐藏层→logits）	dist.gather（主卡收集）	下一个 token 预测（采样）

all_reduce vs gather：
词嵌入层需要所有 GPU 都得到完整嵌入向量（后续层需在每个 GPU 继续计算），因此用 all_reduce 让所有 GPU 汇总结果。
LM Head 仅需主卡得到完整 logits（用于采样输出），其他 GPU 无需，因此用 gather 仅主卡收集，减少通信量。
四、应用场景与依赖
适用场景：
大词汇表 LLM（如词汇表 > 50 万）的分布式训练 / 推理。
单卡内存不足，需通过张量并行（TP）拆分参数的场景（如 70B、175B 参数模型）。
依赖条件：
需提前初始化 torch.distributed 环境（如 dist.init_process_group(backend="nccl", init_method="env://")）。
词汇表大小 num_embeddings 必须能被 tp_size 整除（代码中 assert 强制校验）。
推理时需配合 get_context() 提供的上下文信息（如 is_prefill、cu_seqlens_q），与之前分析的 nano-vllm 推理框架兼容。
总结
这段代码是工业级大模型分布式推理 / 训练的核心组件，通过词汇表分片存储、分片计算、分布式通信汇总三大步骤，解决了大词汇表导致的单卡内存不足问题：

VocabParallelEmbedding 实现并行化词嵌入，确保每个 GPU 仅处理自己的词汇分片，通过 all_reduce 得到完整嵌入。
ParallelLMHead 实现并行化 logits 计算，通过 gather 主卡收集完整 logits，供后续采样使用。

理解这两个类后，可清晰掌握大模型张量并行中 “词汇表维度拆分” 的工程化实现逻辑，以及如何通过分布式通信平衡内存与计算效率。

    """

    def __init__(
        self,
        num_embeddings: int,  # 总词汇表大小（如 100000）
        embedding_dim: int,   # 嵌入维度（如 4096）
    ):
        super().__init__()
        # 1. 分布式参数：tp_rank（当前GPU序号）、tp_size（总GPU数）
        self.tp_rank = dist.get_rank()  # 需提前初始化分布式环境（如 dist.init_process_group）
        self.tp_size = dist.get_world_size()

        # 2. 校验词汇表可均匀分片（避免某GPU多存或少存）
        assert num_embeddings % self.tp_size == 0

        # 3. 计算当前GPU的词汇表分片范围
        self.num_embeddings = num_embeddings
        self.num_embeddings_per_partition = num_embeddings // self.tp_size  # 每GPU存储的词汇数
        self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank  # 分片起始索引
        self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition  # 分片结束索引

        # 4. 定义当前GPU的分片权重（仅存储1/tp_size的词汇嵌入）
        self.weight = nn.Parameter(torch.empty(
            self.num_embeddings_per_partition,
            embedding_dim
        ))

        # 5. 自定义权重加载逻辑（仅加载当前GPU的分片权重）
        self.weight.weight_loader = self.weight_loader

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
        """
        分片权重加载
        核心作用：加载预训练权重时，无需将完整权重（如 16GB）加载到每个 GPU，只需每个 GPU 加载自己的分片（如 4GB），大幅降低内存压力。
        """
        param_data = param.data  # 当前GPU的分片权重数据
        shard_size = param_data.size(0)  # 分片大小（num_embeddings_per_partition）
        start_idx = self.tp_rank * shard_size  # 当前分片在完整权重中的起始索引

        # 1. 从完整权重中切分当前GPU的分片（narrow：沿第0维切分，起始位置start_idx，长度shard_size）
        loaded_weight_shard = loaded_weight.narrow(0, start_idx, shard_size)

        # 2. 校验形状匹配，避免加载错误
        assert param_data.size() == loaded_weight_shard.size()

        # 3. 复制分片权重到当前GPU的参数中
        param_data.copy_(loaded_weight_shard)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        分片计算 + 分布式汇总
        关键逻辑拆解：
mask 筛选：确保每个 GPU 仅处理自己分片内的 token，避免无效计算。
索引调整：将全局 token_id 转为分片内的相对索引（如 2500→0），确保 F.embedding 能正确查找当前 GPU 的分片权重。
分布式汇总（all_reduce）：每个 GPU 的 y 中，仅负责的 token 有非零嵌入，其他为 0；通过 dist.all_reduce 加和所有 GPU 的 y，
得到每个 token 的完整嵌入（非负责 GPU 的 0 不影响结果）。
        """
        # x: 输入token_id，形状通常为 [batch_size, seq_len] 或 [batch_size*seq_len]

        # 1. 若使用词汇表并行（tp_size>1），筛选当前GPU负责的token并调整索引
        if self.tp_size > 1:
            # mask：标记x中属于当前GPU分片的token（True=负责，False=不负责）
            mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
            # 调整token_id为分片内的相对索引（如rank1的2500→0，2501→1）
            x = mask * (x - self.vocab_start_idx)

        # 2. 词嵌入 lookup：仅计算当前GPU分片内的token嵌入（其他token因x=0会取权重第0行，但后续会被mask置零）
        y = F.embedding(x, self.weight)  # 形状：[batch_size, seq_len, embedding_dim]

        # 3. 若使用并行，掩码无效嵌入 + 分布式汇总（all_reduce）
        if self.tp_size > 1:
            # 扩展mask到embedding_dim维度（避免影响嵌入向量的每个维度）
            mask_expanded = mask.unsqueeze(1)  # 形状：[batch_size, 1, seq_len]（适配y的维度）
            y = mask_expanded * y  # 无效token的嵌入置为0

            # 所有GPU汇总结果：每个token的嵌入仅由负责它的GPU贡献非零值，其他GPU贡献0，加和后得到完整嵌入
            dist.all_reduce(y, op=dist.ReduceOp.SUM)

        return y

"""
二、ParallelLMHead：并行化语言模型输出头
继承自 VocabParallelEmbedding，负责将模型的隐藏层输出映射为词汇表维度的 logits（用于预测下一个 token），核心是分片线性变换 + 分布式收集完整 logits。
"""


class ParallelLMHead(VocabParallelEmbedding):
    def __init__(
            self,
            num_embeddings: int,
            embedding_dim: int,
            bias: bool = False,  # 是否使用偏置项
    ):
        """
        扩展 bias 支持
        bias 分片：若使用偏置项，每个 GPU 仅存储自己分片词汇表对应的 bias（与权重分片大小一致），避免全量 bias 占用内存。
        """
        super().__init__(num_embeddings, embedding_dim)

        # 1. 若使用bias，定义分片bias参数（与权重分片逻辑一致）
        if bias:
            self.bias = nn.Parameter(torch.empty(self.num_embeddings_per_partition))
            self.bias.weight_loader = self.weight_loader  # 复用分片权重加载逻辑
        else:
            self.register_parameter("bias", None)  # 无bias则注册为None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        分片 logits 计算 + 分布式收集
        核心逻辑拆解：
Prefill 阶段处理：Prefill 是处理输入提示词，模型需计算每个位置的隐藏层，但 LMHead 仅需最后一个位置的隐藏层（用于预测第一个输出 token），因此通过 last_indices 提取最后一个 token 的隐藏层。
分片线性变换：F.linear(x, self.weight, self.bias) 将隐藏层（embedding_dim）映射到当前 GPU 负责的分片词汇表维度（num_embeddings_per_partition），得到分片 logits。
分布式收集（gather）：
仅 rank0（主卡）收集所有 GPU 的分片 logits，其他 GPU 无需完整 logits（推理时仅主卡需采样输出）。
torch.cat(all_logits, dim=-1) 将分片 logits 沿词汇表维度拼接，得到完整 logits（形状：[batch_size, num_embeddings]），供后续采样（如 argmax、sample）使用。
        """
        # x: 模型隐藏层输出，形状：[batch_size*seq_len, embedding_dim]（prefill阶段）或 [batch_size, embedding_dim]（decode阶段）

        # 1. 获取推理上下文，处理prefill阶段（仅需序列最后一个token的logits）
        context = get_context()  # 复用之前的推理上下文（含is_prefill、cu_seqlens_q等）
        if context.is_prefill:
            # cu_seqlens_q：累计序列长度（如 [0, 3, 8] 表示2个序列，长度3和5）
            # last_indices：每个序列的最后一个token的索引（3-1=2，8-1=7）
            last_indices = context.cu_seqlens_q[1:] - 1
            x = x[last_indices].contiguous()  # 取每个序列最后一个token的隐藏层，确保内存连续

        # 2. 线性变换：将隐藏层映射到分片词汇表维度（logits分片）
        # 注：self.weight 是词嵌入层的权重（通常与Embedding层共享权重，此处未显式写共享，需外部处理）
        logits = F.linear(x, self.weight, self.bias)  # 形状：[batch_size, num_embeddings_per_partition]

        # 3. 若使用并行，收集所有GPU的分片logits，拼接为完整logits
        if self.tp_size > 1:
            # rank0 作为收集者：创建空列表存储所有GPU的分片logits
            all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None

            # 分布式收集：所有GPU将自己的logits发送给rank0，rank0按顺序存储到all_logits
            dist.gather(logits, all_logits, dst=0)

            # rank0 拼接所有分片logits（沿词汇表维度，-1），其他rank设为None（仅主卡需完整logits用于采样）
            logits = torch.cat(all_logits, dim=-1) if self.tp_rank == 0 else None

        return logits
